import requests
from typing import Dict, Any, Optional
from core.plugin_base import PluginBase
from core.data_loader import DataLoader

class APIPlugin(PluginBase):
    """支持REST API数据源的插件，支持动态注册到DataLoader"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.supported_apis = {
            'rest': self._load_rest_api,
            'graphql': self._load_graphql,
            'coze_digital_human': self._load_coze_digital_human  # 集成Coze数字人API[2](@ref)
        }
        self.auth_types = ['none', 'api_key', 'oauth']
        
    def register(self, loader: DataLoader):
        """注册插件到数据加载器"""
        for api_type in self.supported_apis:
            loader.supported_sources[api_type] = self.supported_apis[api_type]
    
    def _load_rest_api(self, url: str, **kwargs) -> Dict[str, Any]:
        """加载REST API数据"""
        params = kwargs.get('params', {})
        headers = self._generate_headers(kwargs.get('auth'))
        
        try:
            response = requests.get(url, params=params, headers=headers)
            response.raise_for_status()
            return self._normalize_data(response.json())
        except Exception as e:
            raise ValueError(f"REST API请求失败: {e}")
    
    def _load_graphql(self, url: str, **kwargs) -> Dict[str, Any]:
        """加载GraphQL数据"""
        query = kwargs.get('query')
        variables = kwargs.get('variables', {})
        headers = self._generate_headers(kwargs.get('auth'))
        
        try:
            response = requests.post(
                url, 
                json={'query': query, 'variables': variables},
                headers=headers
            )
            return self._normalize_data(response.json())
        except Exception as e:
            raise ValueError(f"GraphQL请求失败: {e}")
    
    def _load_coze_digital_human(self, url: str, **kwargs) -> Dict[str, Any]:
        """集成Coze数字人API（示例）[2](@ref)"""
        action = kwargs.get('action')
        access_token = kwargs.get('auth').get('api_key')
        
        endpoints = {
            'create_video': f"{url}/create_video_by_text",
            'get_video': f"{url}/get_video"
        }
        
        if action == 'create':
            data = {
                'text': kwargs.get('text'),
                'voice_type': kwargs.get('voice_type', 'default')
            }
            response = requests.post(
                endpoints['create_video'],
                headers={'Authorization': f'Bearer {access_token}'},
                json=data
            )
            return {'task_id': response.json().get('task_id')}
        
        elif action == 'query':
            response = requests.get(
                f"{endpoints['get_video']}?task_id={kwargs.get('task_id')}",
                headers={'Authorization': f'Bearer {access_token}'}
            )
            return response.json()
    
    def _generate_headers(self, auth_config: Optional[Dict]) -> Dict[str, str]:
        """生成API请求头"""
        headers = {'User-Agent': 'MiniBox-Pro/1.0'}
        if not auth_config:
            return headers
            
        if auth_config['type'] == 'api_key':
            if auth_config['position'] == 'header':
                headers[auth_config['key_name']] = auth_config['key_value']
            elif auth_config['position'] == 'query':
                pass  # Query参数由requests自动处理
        elif auth_config['type'] == 'oauth':
            headers['Authorization'] = f"Bearer {auth_config['token']}"
        
        return headers
    
    def _normalize_data(self, raw_data: Any) -> Dict[str, Any]:
        """标准化API返回数据为DataLoader兼容格式"""
        if isinstance(raw_data, list):
            return {'data': raw_data}
        elif isinstance(raw_data, dict):
            return raw_data
        else:
            raise ValueError("API返回数据格式不支持")

# 插件注册入口
def register_plugin(loader: DataLoader) -> APIPlugin:
    config = {
        'name': 'API Plugin',
        'version': '1.0',
        'description': '支持REST/GraphQL/Coze数字人等API数据源'
    }
    plugin = APIPlugin(config)
    plugin.register(loader)
    return plugin